﻿using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using OpenCvSharp;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace Common.plate.common
{
    /// <summary>
    /// YOLOv8 车牌检测模型实现
    /// </summary>
    public class YoloV8 : IDisposable
    {
        private InferenceSession _session; // ONNX推理会话
        private int _inputWidth; // 模型输入宽度
        private int _inputHeight; // 模型输入高度
        private float _confidenceThreshold = 0.5f; // 置信度阈值
        private float _iouThreshold = 0.4f; // IOU阈值（非极大值抑制）



        /// <summary>
        /// 初始化YOLOv8模型
        /// </summary>
        /// <param name="modelPath">ONNX模型路径</param>
        public YoloV8(string modelPath)
        {
            // 加载ONNX模型
            var options = new SessionOptions(){
                LogSeverityLevel = OrtLoggingLevel.ORT_LOGGING_LEVEL_WARNING // 只输出警告日志，减少冗余
            };
            _session = new InferenceSession(modelPath, options);

            // 获取模型输入尺寸（假设输入形状为[1, 3, H, W]）
            var inputMeta = _session.InputMetadata.Values.First();
            var shape = inputMeta.Dimensions;
            _inputHeight = shape[2];//(640)
            _inputWidth = shape[3];//(640)
        }

        /// <summary>
        /// 检测车牌
        /// </summary>
        /// <param name="image">输入图片</param>
        /// <returns>检测结果列表</returns>
        public List<DetectionResult> Detect(Mat img)
        {
            // 获取模型输入节点名称
            var inputName = _session.InputNames[0];
            // 1. 图像预处理
            var (inputTensor, scale, pad) = Preprocess(img);

            // 2. 模型推理
            var inputs = new List<NamedOnnxValue>
            {
            NamedOnnxValue.CreateFromTensor(inputName, inputTensor) // 输入名称需与模型一致
            };
            // 执行推理，获取输出结果
            using var outputs = _session.Run(inputs);
            var outputTensor = outputs.First().AsTensor<float>();

            // 3. 解析输出并后处理（非极大值抑制）(筛选车牌框)
            var results = Postprocess(outputTensor, img.Cols, img.Rows, scale, pad);
            return results;
        }

        /// <summary>
        /// 图像预处理：缩放、归一化、转换为张量
        /// </summary>
        private (DenseTensor<float> tensor, float scale, (int w, int h) pad) Preprocess(Mat img)
        {
            // 保存原始尺寸
            int originalWidth = img.Cols;
            int originalHeight = img.Rows;

            // 计算缩放比例（保持宽高比）
            float scale = Math.Min((float)_inputWidth / originalWidth, (float)_inputHeight / originalHeight);
            int newWidth = (int)(originalWidth * scale);
            int newHeight = (int)(originalHeight * scale);
            // 重新计算以确保尺寸准确
            newWidth = Math.Min(newWidth, _inputWidth);
            newHeight = Math.Min(newHeight, _inputHeight);

            // 缩放图像
            Mat resized = new Mat();
            Cv2.Resize(img, resized, new Size(newWidth, newHeight));
            // 计算填充（使图像尺寸符合模型输入，避免拉伸）
            int padW = (_inputWidth - newWidth) / 2;
            int padH = (_inputHeight - newHeight) / 2;

            

            // 填充（上下左右填充黑色）
            Mat padded = new Mat();
            Cv2.CopyMakeBorder(resized, padded, padH, _inputHeight - newHeight - padH,
                              padW, _inputWidth - newWidth - padW, BorderTypes.Constant, Scalar.Black);

            // 转换为RGB（YOLO通常使用RGB输入，而OpenCV默认是BGR）
            Mat rgb = new Mat();
            Cv2.CvtColor(padded, rgb, ColorConversionCodes.BGR2RGB);

            // 归一化（除以255，转换为float）
            rgb.ConvertTo(rgb, MatType.CV_32FC3, 1.0 / 255.0);

            // 转换为张量（形状：[1, 3, H, W]）
            var tensor = new DenseTensor<float>(new[] { 1, 3, _inputHeight, _inputWidth });
            for (int c = 0; c < 3; c++)
            {
                for (int y = 0; y < _inputHeight; y++)
                {
                    for (int x = 0; x < _inputWidth; x++)
                    {
                        tensor[0, c, y, x] = (float)rgb.At<Vec3f>(y, x)[c];
                    }
                }
            }

            return (tensor, scale, (padW, padH));
        }

        /// <summary>
        /// 后处理：解析输出张量，应用非极大值抑制（NMS）
        /// </summary>
        private List<DetectionResult> Postprocess(Tensor<float> output, int originalWidth, int originalHeight,
                                                 float scale, (int w, int h) pad)
        {
            var results = new List<DetectionResult>();
            int numClasses = 1; // 假设只检测"车牌"这一类

            // YOLOv8输出格式：[batch,4 + 1 + num_classes, num_boxes, ]
            // 其中4是边界框（x_center, y_center, width, height），1是置信度，剩余是类别概率
            int numBoxes = output.Dimensions[1];//候选框数
            for (int i = 0; i < numBoxes; i++)
            {
                // 解析边界框（中心坐标、宽高）
                float xCenter = output[0,i, 0];//x坐标
                float yCenter = output[0, i, 1];//y坐标
                float width = output[0, i, 2];//宽
                float height = output[0,i, 3];//高

                // 解析置信度和类别
                float confidence = output[0, i, 4];
                if (confidence < _confidenceThreshold)
                    continue; // 过滤低置信度

                // 转换边界框为左上角和右下角坐标（相对于模型输入尺寸）
                float x1 = xCenter - width / 2;
                float y1 = yCenter - height / 2;
                float x2 = xCenter + width / 2;
                float y2 = yCenter + height / 2;

                // 还原到原始图像尺寸（去除填充和缩放）
                x1 = (x1 - (float)pad.w) / scale;
                y1 = (y1 - (float)pad.h) / scale;
                x2 = (x2 - (float)pad.w) / scale;
                y2 = (y2 - (float)pad.h) / scale;

                // 确保坐标在图像范围内
                x1 = Math.Max(0, Math.Min(x1, originalWidth));  // x1 ≥0 且 ≤原始宽
                y1 = Math.Max(0, Math.Min(y1, originalHeight)); // y1 ≥0 且 ≤原始高
                x2 = Math.Max(0, Math.Min(x2, originalWidth));  // 同理限制x2、y2
                y2 = Math.Max(0, Math.Min(y2, originalHeight));

                // 4. 过滤无效框（确保坐标在原图范围内）
                if (x2 > x1 && y2 > y1
                    && x1 >= 0 && y1 >= 0 &&
                    x2 <= originalWidth && y2 <= originalHeight)
                {
                    results.Add(new DetectionResult
                    {
                        Box = new Rect((int)x1, (int)y1, (int)(x2 - x1), (int)(y2 - y1)),
                        Confidence = confidence
                    });
                }                
            }
            // 应用非极大值抑制（NMS）去除重叠框
            return NMS(results);
        }

        /// <summary>
        /// 非极大值抑制：去除重叠度高的边界框
        /// </summary>
        private List<DetectionResult> NMS(List<DetectionResult> boxes)
        {
            var result = new List<DetectionResult>();
            if (boxes.Count == 0) return result;

            // 按置信度降序排序
            var sorted = boxes.OrderByDescending(b => b.Confidence).ToList();

            while (sorted.Count > 0)
            {
                var current = sorted[0];
                result.Add(current);
                sorted.RemoveAt(0);

                // 移除与当前框IOU大于阈值的框
                sorted.RemoveAll(b => CalculateIOU(current.Box, b.Box) > _iouThreshold);
            }

            return result;
        }

        /// <summary>
        /// 计算两个边界框的交并比（IOU）
        /// </summary>
        private float CalculateIOU(Rect a, Rect b)
        {
            int areaA = a.Width * a.Height;
            int areaB = b.Width * b.Height;

            int x1 = Math.Max(a.Left, b.Left);
            int y1 = Math.Max(a.Top, b.Top);
            int x2 = Math.Min(a.Right, b.Right);
            int y2 = Math.Min(a.Bottom, b.Bottom);

            int intersection = Math.Max(0, x2 - x1) * Math.Max(0, y2 - y1);
            return (float)intersection / (areaA + areaB - intersection);
        }

        /// <summary>
        /// YOLO检测结果（边界框和置信度）
        /// </summary>
        public class DetectionResult
        {
            public Rect Box { get; set; } // 边界框（左上角x,y,宽,高）
            public float Confidence { get; set; } // 置信度
        }

        public void Dispose()
        {
            _session.Dispose();
        }
    }
}
